# -*- coding: utf-8 -*-
"""
image resize.
"""

import cv2
import sys
import argparse
import numpy as np
import voc_xml
from voc_xml import CreateXML
import os


def resize_xy(x, y, fx, fy):
    """
    放缩点坐标
    Args:
        x,y:待放缩点坐标
        fx,fy:放缩比例
    return:
        x,y:放缩后坐标点
    """
    return int(x*fx), int(y*fy)


def resize_box(box, fx, fy, left, top):
    """
    放缩目标框：
    Args:
        box: 目标框 [xmin,ymin,xmax,ymax]
        fx,fy: x,y坐标轴放缩比例
    return:
        rsize_box: 放缩后的坐标框 [xmin,ymin,xmax,ymax]
    """
    xmin, ymin = resize_xy(box[0], box[1], fx, fy)
    xmin, ymin = xmin + left, ymin + top
    xmax, ymax = resize_xy(box[2], box[3], fx, fy)
    xmax, ymax = xmax + left, ymax + top
    return [xmin, ymin, xmax, ymax]


def resize_img(src, dsize=(0, 0), fx=1.0, fy=1.0, fixed_ratio=False):
    """
    放缩图片
    Args:
        src:源图片
        dsize:指定放缩大小(w,h)
        fx,fy:比例放缩
    return：
        sized_img:放缩后的图像
    """
    resized_w = dsize[0]
    resized_h = dsize[1]
    resized_img = np.zeros((resized_h, resized_w, 3), dtype=np.uint8) + 255
    src_h, src_w, _ = src.shape
    if fixed_ratio:
        if resized_w / src_w < resized_h / src_h:
            new_w = resized_w
            new_h = src_h * (resized_w / src_w)
            if new_h > resized_h:
                left, top = int(0), int((new_h - resized_h) / 2)
            else:
                left, top = int(0), int((resized_h - new_h) / 2)
        else:
            new_w = src_w * (resized_h / src_h)
            new_h = resized_h
            if new_w > resized_w:
                left, top = int((new_w - resized_w) / 2), int(0)
            else:
                left, top = int((resized_w - new_w) / 2), int(0)
    else:
        new_w = resized_w
        new_h = resized_h
        left, top = 0, 0
    sized_img = cv2.resize(src, (int(new_w), int(new_h)), fx=fx, fy=fy)
    resized_img[top:top+sized_img.shape[0], left:left+sized_img.shape[1]] = sized_img
    return resized_img, left, top


def resize_xml(resized_img_name, xml_tree, dsize=(0, 0),
               fx=1.0, fy=1.0, fixed_ratio=False):
    """
    xml目标框放缩变换
    Args:
        resized_img_name: resize图片保存名
        xml_tree:  待resize xml  ET.parse()
        dsize:指定放缩大小(w,h)
        fx,fy:比例放缩
    return:
        createdxml : 创建的xml CreateXML对象        
    """
    root = xml_tree.getroot()
    size = root.find('size')
    src_w, src_h, depth = int(size.find('width').text),\
                          int(size.find('height').text),\
                          int(size.find('depth').text)
    resize_imgw, resize_imgh = src_w, src_h
    if dsize == (0, 0):
        resized_w = int(src_w * fx)
        resized_h = int(src_h * fy)
    else:
        resized_w, resized_h = dsize

    if fixed_ratio:
        if resized_w / src_w < resized_h / src_h:
            new_w = resized_w
            new_h = src_h * (resized_w / src_w)
            resize_fx, resize_fy = resized_w / src_w, resized_w / src_w
            if new_h > resized_h:
                left, top = int(0), int((new_h - resized_h) / 2)
            else:
                left, top = int(0), int((resized_h - new_h) / 2)
        else:
            new_w = src_w * (resized_h / src_h)
            new_h = src_h
            resize_fx, resize_fy = resized_h / src_h, resized_h / src_h
            if new_w > resized_w:
                left, top = int((new_w - resized_w) / 2), int(0)
            else:
                left, top = int((resized_w - new_w) / 2), int(0)
    else:
        new_w = resized_w
        new_h = resized_h
        left, top = 0, 0
        resize_fx, resize_fy = resized_w/src_w, resized_h/src_h

    createdxml = CreateXML(resized_img_name, resized_w, resized_h, depth)
    
    for obj in root.iter('object'):
        obj_name = obj.find('name').text
        xml_box = obj.find('bndbox')
        xmin = int(xml_box.find('xmin').text)
        ymin = int(xml_box.find('ymin').text)
        xmax = int(xml_box.find('xmax').text)
        ymax = int(xml_box.find('ymax').text)
        box = resize_box([xmin, ymin, xmax, ymax], resize_fx, resize_fy, left, top)
        if (box[0] >= box[2]) or (box[1] >= box[3]):
            continue
        createdxml.add_object_node(obj_name, box[0], box[1], box[2], box[3])
    return createdxml


def generate_resizeImg_xml(img, xml_tree, resized_img_name,
                           dsize=(0, 0), fx=1.0, fy=1.0, fixed_ratio=False):
    """
    生成旋转后的图片和xml文件
    Args:
        img:源图片
        xml_tree：待resizexml  ET.parse()
        resized_img_name: resize图片保存名
        dsize:指定放缩大小(w,h)
        fx,fy:比例放缩
    return:
        resized_img,resized_xml       
    """
    resized_img, _, _ = resize_img(img, dsize, fx, fy, fixed_ratio)
    resized_xml = resize_xml(resized_img_name, xml_tree, dsize, fx, fy, fixed_ratio)
    return resized_img, resized_xml


def batch_resize(imgs_dir, xmls_dir, imgs_save_dir, xmls_save_dir, img_suffix,
                 name_suffix, new_size=(0, 0), scale_w=1.0, scale_h=1.0, fixed_ratio=False):
    """
    放缩指定路径下的图片和xml
    Args:
        imgs_dir,xmls_dir: 待放缩图片、原始xml文件存储路径
        imgs_save_dir，xmls_save_dir: 处理完成的图片、xml文件存储路径
        img_suffix: 图片可能的后缀名['.jpg','.png','.bmp',..]
        name_suffix: 处理完成的图片、xml的命名标识
        dsize: 指定放缩大小(w,h)
        fx,fy: 比例放缩       
    """
    finished_num = 0
    for root, dirs, files in os.walk(xmls_dir):
        img_num = len(files)
        for xml_name in files:
            xml_file = os.path.join(xmls_dir, xml_name)
            img_file = None
            for suffix in img_suffix:
                if os.path.exists(os.path.join(imgs_dir, xml_name.split('.')[0]+suffix)):
                    img_file = os.path.join(imgs_dir, xml_name.split('.')[0]+suffix)
                    break
            if img_file is None:
                print("there has no image for ", xml_name)
                continue
            img = cv2.imread(img_file)
            
            imgh,imgw,n_channels = img.shape
            resize_imgw,resize_imgh = imgw,imgh
            if new_size == (0, 0):
                resize_imgw = imgw * scale_w
                resize_imgh = imgh * scale_h
            else:
                resize_imgw, resize_imgh = new_size

            resized_img_name = "{}_{}_w{}_h{}.{}".format(xml_name.split('.')[0], name_suffix,
                                                         resize_imgw, resize_imgh, img_file.split('.')[-1])
            img_resize, xml_resize = generate_resizeImg_xml(img, voc_xml.get_xml_tree(xml_file),
                                                            resized_img_name, new_size, scale_w,
                                                            scale_h, fixed_ratio)
            cv2.imwrite(os.path.join(imgs_save_dir, resized_img_name), img_resize)
            xml_resize.save_xml(xmls_save_dir, resized_img_name.split('.')[0]+'.xml')
            finished_num += 1
            if finished_num % 100 == 0:
                print("Resizing: {}th image, total: {}".format(finished_num, img_num))
        print('*'*80)
        print()
        print("finished: {} images, total: {}".format(finished_num, img_num))


def start_resize(root_dir, save_dir):
    """

    :param root_dir:
    :return:
    """
    img_dir = os.path.join(root_dir, "JPEGImages")
    xml_dir = os.path.join(root_dir, "Annotations")
    resize_img_save_dir = os.path.join(save_dir, "JPEGImages")
    resize_xml_save_dir = os.path.join(save_dir, "Annotations")
    if not os.path.exists(resize_img_save_dir):
        os.makedirs(resize_img_save_dir)
    if not os.path.exists(resize_xml_save_dir):
        os.makedirs(resize_xml_save_dir)
    img_suffix = ['.jpg', '.png', '.bmp']
    # 命名标识
    name_suffix = 'rsize'

    # 指定放缩大小(w,h)
    new_size = (960, 540)

    # 放缩比例
    scale_w = 1.0
    scale_h = 1.0

    # 是否保持宽高比。若保持宽高比则不足部分填充黑边
    fixed_ratio = True

    batch_resize(img_dir, xml_dir, resize_img_save_dir, resize_xml_save_dir,
                           img_suffix, name_suffix, new_size, scale_w, scale_h, fixed_ratio)
    

if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        SAVE_DIR = sys.argv[2]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                            help="VOC dataset root directory")
        PARSER.add_argument("--save_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/resized_voc"),
                            help="result save dir")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        SAVE_DIR = ARGS.save_dir
    start_resize(ROOT_DIR, SAVE_DIR)














